import matplotlib.pyplot as plt

from pathlib import Path
import numpy as np
import cv2
from vae.mnist_dataset import MNISTDataset
from vae.vae_model import VAE, VAELoss
from torch.utils.data import DataLoader
import torch
from torch import optim

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

base_path = Path("D:\\datasets\\image\\MNIST\\mnist_dataset")

sets = {
    "train": MNISTDataset(base_path / "train", base_path / "train_labs.txt", True),
    "test": MNISTDataset(base_path / "test", base_path / "test_labs.txt", True)
    }


loaders = {key: DataLoader(sets[key], batch_size=128, shuffle=True) for key in sets}

model = VAE(1, 256, 16, 784).to(device)
loss_fn = VAELoss()
optimizer = optim.Adagrad(model.parameters(), lr=1e-3, weight_decay=1e-4)

epoch_num = 1000


best_loss = 1e+19
for i in range(epoch_num):
    
    for mode in loaders:
        loss_val = 0.0
        with torch.set_grad_enabled(mode=="train"):

            for inputs, labels in loaders[mode]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                mean, var, outputs = model(inputs)

                loss = loss_fn(mean, var, outputs, labels)

                loss_val += loss.item()

                if mode == "train":
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()

            
            print("mode:{} epoch:[{}/{}] lossVal:{:.4f}".format(mode, i + 1, epoch_num, loss_val / len(sets[mode])))

            if loss_val < best_loss:
                best_loss = loss_val

                torch.save(model.state_dict(), "./best_model")